from math import log,inf
from HiddenMarkovModel import Print_Matrix

class Matrix(object):
    def __init__(self,data:dict[str,dict]) -> None:
        self.row_name = tuple(data.keys())
        self.colum_name = tuple(data[self.row_name[0]].keys())
        self.mat = data

    def copy_empty(self) -> "Matrix":
        data = {r:{c:None for c in self.colum_name}
                for r in self.row_name}
        return Matrix(data)
    
    def log(self,base=2) -> "Matrix":
        log_mat = self.copy_empty()
        for i in self.row_name:
            for j in self.colum_name:
                if self[i,j] == 0:
                    log_mat[i,j] = -inf
                else:
                    log_mat[i,j] = log(self[i,j],base)
        return log_mat

    def __getitem__(self,keys) -> float|tuple:
        if isinstance(keys,str):
            if keys in self.row_name:
                return tuple(self.mat[keys].values())
            else:
                raise ValueError("error key")
        elif isinstance(keys,tuple):
            i,j = keys
            if isinstance(i,slice):
                if j in self.colum_name:
                    return tuple([col[j] for col in self.mat.values()])
                else:
                    raise ValueError("error key")
            elif i in self.row_name and j in self.colum_name:
                return self.mat[i][j]
            else:
                raise ValueError("error key")
        else:
            raise ValueError("error key type")
    
    def __setitem__(self,keys:tuple,value) -> None:
        if not isinstance(keys,tuple):
            raise ValueError("error key type")
        
        i,j = keys
        if isinstance(i,slice):
            raise ValueError("error key type")
            
        if i in self.row_name and j in self.colum_name:
            self.mat[i][j] = value
        else:
            raise ValueError("error key")
    
    def keys(self) -> tuple[tuple]:
        return (self.row_name,self.colum_name)
    
    def View(self):
        Print_Matrix(self.mat)

def test_mat():
    A = {"5" : {"A":0.20, "C":0.30, "G":0.30, "T":0.20},
         "M" : {"A":0.25, "C":0.25, "G":0.25, "T":0.25},
         "3" : {"A":0.35, "C":0.15, "G":0.15, "T":0.35}}
    A = Matrix(A)
    print(A[:,"C"])
    print(A["M"])
    A["5","C"] = 10
    print(A[:,"C"])
    A.View()


class Node(object):
    def __init__(self,state,step) -> None:
        self.state = state
        self.step = step
        self.value = None
        self.Next = []
    
    def set_value(self,v) -> None:
        self.value = v

    def add_next(self,node:"Node",weight=None) -> None:
        self.Next.append([node,weight])

class Network(object):
    def __init__(self,state,observed_seq) -> None:
        self.states = state
        self.obs = observed_seq
        
        Network = {i:{s:Node(s,i) for s in self.states}
                   for i in range(len(self.obs))}
        self.N = Matrix(Network)
    
    def __getitem__(self,keys) -> "Node":
        return self.N[keys]
    
    def add_edge(self,node_from:tuple,node_to:tuple,weight=None):
        node_from = self.N[node_from]
        node_to = self.N[node_to]
        node_from.add_next(node_to,weight)

    def all_path(self,start:Node):
        if start.step == len(self.obs)-1:
            return [start.state],[start.value]
        if not start.Next:
            return False
        
        paths, values = [], []
        for next,t_v in start.Next:
            vs = self.all_path(next)
            if not vs:
                continue
            paths.extend([start.state+path for path in vs[0]])
            values.extend([start.value+t_v+value for value in vs[1]])
        
        return paths,values

def test_network():
    N = Network(("A","B","C"),"ASDSADASfdf")
    print(N[1,"A"].step)
    print(N[8,"C"].state)
    N.add_edge((1,"A"),(2,"B"),4)
    print(N[1,"A"].Next)
    print(N[1,"A"].Next[0][0].step)

if __name__ == "__main__":
    # test_mat()
    test_network()
    pass